import pandas as pd

# 生成对应数据集内所有一项集
def createOneitmeset(dataSet):
    one_itmeset = []
    # 遍历数据集
    for transaction in dataSet:
        for item in transaction:
            # 将未出现的项存储在一项集中
            if not [item] in one_itmeset:
                one_itmeset.append([item])
    one_itmeset.sort()
    return list(map(frozenset, one_itmeset))

# 生成k项集
def K_itemsetGenerator(K_itemset_prev, K):
    # K-项集
    K_itemset = []
    # 前一项集项数
    len_prev = len(K_itemset_prev)
    # 二重循环两两组合生成生成新项集
    for i in range(len_prev):
        for j in range(i+1, len_prev):
            L1=list(K_itemset_prev[i])[:K-2]
            L2=list(K_itemset_prev[j])[:K-2]
            L1.sort()
            L2.sort()
            if L1==L2:
                K_itemset.append(K_itemset_prev[i]|K_itemset_prev[j])
    return K_itemset

# 一个项集是非频繁集，那么它的所有超集（包含该非频繁集的父集）也是非频繁的）来进行滤除
def filter_itmeset(Data, itmeset):
    item_frequency = {}
    # 循环遍历项集中元素，并记录出现次数至item_frequency字典中
    for item in Data:
        for element in itmeset:
            # 判读是否为子集
            if element.issubset(item):
                if not element in item_frequency:
                    item_frequency[element] = 1
                else:
                    item_frequency[element] += 1

    # 数据集中总的项集个数
    numItems = float(len(Data))
    # 频繁项集
    item_filtered = []
    # 项集中各项支持度
    itemSupprotData={}
    for key in item_frequency:
        # 计算支持度
        support = item_frequency[key]/numItems
        # 满足最小支持度添加到频繁项集列表中
        if support>=minSupport:
            item_filtered.insert(0, key)
        itemSupprotData[key] = support
    return item_filtered, itemSupprotData

def itemset_exhaustion(dataSet):
    one_itmeset = createOneitmeset(dataSet)
    Data = list(map(set, dataSet))
    one_itmeset_filtered, itemSupprotData_union = filter_itmeset(Data, one_itmeset)
    itemset_union = [one_itmeset_filtered]
    K = 2
    while(len(itemset_union[K - 2]) > 0):
        K_itemset = K_itemsetGenerator(itemset_union[K-2], K)
        K_itemset_filtered, K_itemSupprotData = filter_itmeset(Data, K_itemset)
        itemSupprotData_union.update(K_itemSupprotData)
        itemset_union.append(K_itemset_filtered)
        K += 1
    return itemset_union, itemSupprotData_union

def findRelatedItemGroup(K_itemset_freq,itemSupprotData):
    RelatedItemGroup = []
    K = len(K_itemset_freq)
    for i in range(1, K):
        for itemset in K_itemset_freq[i]:
            itemWaitjudge = [frozenset([element]) for element in itemset]
            if(i > 1):
                reunionAcalcConf(itemset,itemWaitjudge,itemSupprotData,RelatedItemGroup)
            else:
                calcConf(itemset,itemWaitjudge,itemSupprotData,RelatedItemGroup)
    return RelatedItemGroup

def calcConf(itemset, itemWaitjudge, itemSupprotData, RelatedItemGroup):
    # 存储itemWaitjudge项中强关联规则
    item_realted = []
    for item in itemWaitjudge:
        # 计算置信度
        conf = itemSupprotData[itemset] / itemSupprotData[itemset - item]
        # 满足条件 存储并输出
        if conf >= minConf:
            print (itemset-item,'===>',item,' confidence:',conf)
            RelatedItemGroup.append((itemset - item, item, conf))
            item_realted.append(item)
    return item_realted

def reunionAcalcConf(itemset, itemWaitjudge, itemSupprotData, RelatedItemGroup):
    # K:当前项集的项数
    K = len(itemWaitjudge[0])
    # 当当前项集项数已经组合到(父项集项数-1)时无需再次递归组合，否则继续在组合
    if (len(itemset) > (K + 1)):
        # 项集包含项再组合，项数为当前项数 + 1
        itemset_reunion = K_itemsetGenerator(itemWaitjudge, K + 1)
        # 计算置信度并进行滤除
        itemset_reunion_filtered = calcConf(itemset, itemset_reunion, itemSupprotData, RelatedItemGroup)
        # 如果没有强关联规则存在，后续的也无须递归再组合计算置信度
        if(len(itemset_reunion_filtered) > 1):
            # 存在强关联规则，递归在组合计算置信度
            reunionAcalcConf(itemset, itemset_reunion_filtered, itemSupprotData, RelatedItemGroup)

# 最小支持度
minSupport = 0.025
# 最小置信度
minConf = 0.5
movies = pd.read_csv('movies.csv')
pd.set_option('display.width', None)
print('数据集如下：')
print(movies)
tempDataSet = movies['genres'].tolist()
dataSet = []
for item in tempDataSet:
    dataSet.append(item.split('|'))
L,supportData=itemset_exhaustion(dataSet)
rules = findRelatedItemGroup(L,supportData)